import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.keras.callbacks import ReduceLROnPlateau

from model import create_generic_model
from data_utils import preprocess_data, normalize_data, create_dataset, load_dataset_from_files, preprocess_dataset, augment_data
from framework import CNNFramework
from nnom import * # type: ignore

# 导入配置参数
from config import *

# 构建模型并打印摘要
model = create_generic_model(INPUT_SHAPE, NUM_CLASSES)
model.summary()

# 创建框架实例
framework = CNNFramework(model=model, input_shape=INPUT_SHAPE, num_classes=NUM_CLASSES,pretrained_model_path = PRETRAINED_MODEL_PATH)

# 数据路径和参数
train_data_path = TRAIN_DATA_PATH
train_labels_path = TRAIN_LABELS_PATH

# 预处理：加载、归一化、增强、创建 Dataset
# 尝试使用新的数据集加载方式
file_list, labels = load_dataset_from_files(
	train_data_path, MOTION_TO_LABEL, 
	DEF_FILE_FORMAT=DEF_FILE_FORMAT, 
	DEF_FILE_MAX=DEF_FILE_MAX, 
	DEF_USE_COLS=DEF_USE_COLS, 
	DEF_N_ROWS=DEF_N_ROWS
)

# 检查是否成功加载了数据
if len(file_list) == 0:
	raise ValueError("没有加载到任何数据文件，请检查数据路径和文件格式。")

# 预处理数据集
x_train, y_train, x_test, y_test = preprocess_dataset(
	file_list, labels, NUM_CLASSES, input_shape=INPUT_SHAPE
)
print(f"数据预处理完成，训练集大小：{len(x_train)}，测试集大小：{len(x_test)}")
print(f"数据形状：{x_train.shape}")
print(f"标签形状：{y_train.shape}")
# 回调

# 在训练时添加学习率调度器
lr_scheduler = ReduceLROnPlateau(
    monitor='val_loss',    # 监控验证损失
    factor=0.5,            # 学习率衰减因子（变为原来的1/2）
    patience=10,           # 等待10个epoch没有改善
    min_lr=1e-7,          # 最小学习率
    verbose=1              # 打印学习率变化信息
)


callbacks = [EarlyStopping(monitor='val_loss', patience=10), 
			 ModelCheckpoint(DEF_MODEL_NAME, monitor='val_accuracy', save_best_only=True, mode='max'),
             lr_scheduler  
            ]

# 训练
framework.train_model(x_train, y_train, 
					  val_ds=(x_test, y_test), 
					  epochs=EPOCHS, 
					  callbacks=callbacks
					  )

x_test_sample = x_test[:100]  # 使用前100个样本作为校准数据集
y_test_sample = y_test[:100]
test_ds = create_dataset(x_test_sample, y_test_sample, NUM_CLASSES, batch_size=BATCH_SIZE, shuffle=False)

generate_model(model, x_test_sample, format='hwc', name=DEF_MODEL_H_NAME)



# 示例：使用测试数据的一部分进行预测
sample_data = x_test[:5]  # 取5个测试样本
# 数据预处理
new_data = normalize_data(sample_data, method='standard')
predicted_classes, probabilities, all_probabilities = CNNFramework.predict_new_data(model, sample_data)

print("\n预测结果示例:")
for i in range(len(predicted_classes)):
    # 将类别索引转换为标签名称
    class_names = list(MOTION_TO_LABEL.keys())
    predicted_label = class_names[predicted_classes[i]]
    true_label = class_names[np.argmax(y_test[i])]
    
    print(f"样本 {i+1}:")
    print(f"  预测类别: {predicted_label} (概率: {probabilities[i]:.4f})")
    print(f"  真实类别: {true_label}")
    print(f"  所有类别概率: {all_probabilities[i]}")
    print()
    

# 评估与绘图
framework.evaluate_model(test_ds)
framework.plot_training_history()